# !/usr/bin/env python
# -*-coding:utf-8 -*-

"""
# File       : kafka_client.py
# Time       ：2025/5/29 14:36
# Description：
"""
import base64
import json
import os

from utils.kafka_util.kafka_msg_model import MESSAGE_TYPES, OPERATIONS
from utils.log import log
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
from confluent_kafka import Producer, Consumer
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)

class KafkaProducer:
    def __init__(self, bootstrap = "zone", consumer_subscribe = None, max_retries: int = 3):
        self.conf = {
            'bootstrap.servers': os.getenv("ZONE_KAFKA_BOOTSTRAP_SERVERS", "10.141.194.154:29092,10.141.194.155:29092,10.141.194.156:29092") if bootstrap == "zone"
            else os.getenv("UNION_KAFKA_BOOTSTRAP_SERVERS", "10.141.194.84:29092,10.141.194.85:29092,10.141.194.86:29092"),
            'queue.buffering.max.messages': 100000,
            'queue.buffering.max.ms': 100,
            'compression.type': 'gzip',
            'acks': 'all',
            'retries': max_retries,
            'retry.backoff.ms': 200,
        }
        self.producer = None
        self.consumer = None
        self.consumer_subscribe = consumer_subscribe

    def start(self):
        if self.producer is None:
            try:
                self.producer = Producer(self.conf)
                log.info("Kafka producer initialized")
            except Exception as e:
                log.error(f"Failed to initialize Kafka producer: {str(e)}")
                raise

    def delivery_report(self, err, msg):
        """Callback for message delivery."""
        if err is not None:
            log.error(f"Message delivery failed: {err}")
        else:
            pass


    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        retry=retry_if_exception_type(Exception)
    )
    async def send_message(self, topic: str, value: str, key: str = None):
        """Send message to Kafka topic synchronously."""
        if self.producer is None:
            self.start()
        try:
            self.producer.produce(
                topic=topic,
                value=value.encode('utf-8'),
                key=key,
                callback=self.delivery_report
            )
            self.producer.poll(0)  # Trigger callbacks
            self.producer.flush(timeout=5.0)  # Wait for message to be sent
            log.info(f"Sent Kafka message to {topic}: {value}")
        except Exception as e:
            log.error(f"Failed to send message to {topic}: {str(e)}")
            raise


    def consume_message(self, group_id: str = "default_consumer_group"):
        """Consume messages from vm_ecn_changed_notify topic and print them."""
        if self.consumer is None:
            consumer_conf = {
                'bootstrap.servers': self.conf['bootstrap.servers'],
                'group.id': group_id,
                'auto.offset.reset': 'earliest',
                'enable.auto.commit': True
            }
            try:
                self.consumer = Consumer(consumer_conf)
                self.consumer.subscribe(self.consumer_subscribe)
                log.info(f"Kafka consumer initialized for topic {self.consumer_subscribe}")
            except Exception as e:
                log.error(f"Failed to initialize Kafka consumer: {str(e)}")
                raise

        try:
            while True:
                msg = self.consumer.poll(1.0)
                if msg is None:
                    continue
                if msg.error():
                    log.error(f"Consumer error: {msg.error()}")
                    continue
                message_value = msg.value().decode('utf-8') if msg.value() else None
                if msg.topic() in ['vm_ecn_changed_notify', 'scheduler_invoke_games']:
                    message_value: dict = json.loads(message_value)
                    decoded_bytes = base64.b64decode(message_value.get("data")).decode('utf-8')
                    parsed_data = json.loads(decoded_bytes)
                    log.info(
                        f'chedule info: messageType({MESSAGE_TYPES.get(message_value.get("messageType"))}), operation({OPERATIONS.get(message_value.get("operation"))}) messageID({message_value.get("messageID")}),\n{parsed_data}')

        except KeyboardInterrupt:
            log.info("Consumer interrupted by user")
        except Exception as e:
            log.error(f"Error consuming message: {str(e)}")
            raise
        finally:
            self.consumer.close()
            self.consumer = None
            log.info("Kafka consumer closed")

    def close(self):
        """Close Kafka producer."""
        if self.producer:
            self.producer.flush()
            log.info("Kafka producer closed")


#

# kafka_producer.consume_message()
